use async_trait::async_trait;
use rustcloud_core::{ServiceResult, ServiceError};
use crate::tracing::{Span, Trace, TraceExporter, TracingProcessor, TracingConfig, SpanKind, SpanStatus};
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::RwLock;
use std::time::{SystemTime, UNIX_EPOCH};

/// 追踪上下文
#[derive(Debug, Clone)]
pub struct TracingContext {
    /// 当前Trace ID
    pub trace_id: Option<String>,
    /// 当前Span ID
    pub span_id: Option<String>,
    /// 是否被采样
    pub sampled: bool,
    /// 额外的上下文信息
    pub baggage: HashMap<String, String>,
}

impl Default for TracingContext {
    fn default() -> Self {
        Self {
            trace_id: None,
            span_id: None,
            sampled: true,
            baggage: HashMap::new(),
        }
    }
}

impl TracingContext {
    /// 创建新的追踪上下文
    pub fn new() -> Self {
        Self::default()
    }

    /// 从现有Span创建上下文
    pub fn from_span(span: &Span) -> Self {
        Self {
            trace_id: Some(span.trace_id.clone()),
            span_id: Some(span.span_id.clone()),
            sampled: true,
            baggage: HashMap::new(),
        }
    }

    /// 设置baggage
    pub fn set_baggage(&mut self, key: &str, value: &str) {
        self.baggage.insert(key.to_string(), value.to_string());
    }

    /// 获取baggage
    pub fn get_baggage(&self, key: &str) -> Option<&String> {
        self.baggage.get(key)
    }

    /// 转换为HTTP头部
    pub fn to_http_headers(&self) -> HashMap<String, String> {
        let mut headers = HashMap::new();
        
        if let Some(trace_id) = &self.trace_id {
            headers.insert("X-Trace-ID".to_string(), trace_id.clone());
        }
        
        if let Some(span_id) = &self.span_id {
            headers.insert("X-Span-ID".to_string(), span_id.clone());
        }
        
        headers.insert("X-Sampled".to_string(), if self.sampled { "1" } else { "0" }.to_string());
        
        // 添加baggage
        for (key, value) in &self.baggage {
            headers.insert(format!("X-Baggage-{}", key), value.clone());
        }
        
        headers
    }

    /// 从HTTP头部创建上下文
    pub fn from_http_headers(headers: &HashMap<String, String>) -> Self {
        let mut context = Self::new();
        
        if let Some(trace_id) = headers.get("X-Trace-ID") {
            context.trace_id = Some(trace_id.clone());
        }
        
        if let Some(span_id) = headers.get("X-Span-ID") {
            context.span_id = Some(span_id.clone());
        }
        
        if let Some(sampled) = headers.get("X-Sampled") {
            context.sampled = sampled == "1";
        }
        
        // 提取baggage
        for (key, value) in headers {
            if key.starts_with("X-Baggage-") {
                let baggage_key = key.strip_prefix("X-Baggage-").unwrap_or(key);
                context.baggage.insert(baggage_key.to_string(), value.clone());
            }
        }
        
        context
    }
}

/// 追踪管理器
pub struct TracingManager {
    service_name: String,
    processor: Arc<TracingProcessor>,
    active_spans: Arc<RwLock<HashMap<String, Span>>>,
    active_traces: Arc<RwLock<HashMap<String, Trace>>>,
    config: TracingConfig,
}

impl TracingManager {
    /// 创建新的追踪管理器
    pub fn new(
        service_name: String,
        config: TracingConfig,
        exporter: Arc<dyn TraceExporter>,
    ) -> Self {
        let processor = Arc::new(TracingProcessor::new(config.clone(), exporter));
        
        Self {
            service_name,
            processor,
            active_spans: Arc::new(RwLock::new(HashMap::new())),
            active_traces: Arc::new(RwLock::new(HashMap::new())),
            config,
        }
    }

    /// 开始新的根Span
    pub async fn start_root_span(&self, operation_name: &str) -> ServiceResult<Span> {
        let span = Span::new_root(operation_name, &self.service_name);
        
        // 创建新的Trace
        let trace = Trace::new(span.clone());
        
        // 存储活跃的Span和Trace
        {
            let mut active_spans = self.active_spans.write().await;
            active_spans.insert(span.span_id.clone(), span.clone());
        }
        
        {
            let mut active_traces = self.active_traces.write().await;
            active_traces.insert(trace.trace_id.clone(), trace);
        }
        
        Ok(span)
    }

    /// 开始新的子Span
    pub async fn start_child_span(
        &self,
        parent: &Span,
        operation_name: &str,
    ) -> ServiceResult<Span> {
        let span = parent.new_child(operation_name, &self.service_name);
        
        // 将子Span添加到对应的Trace
        {
            let mut active_traces = self.active_traces.write().await;
            if let Some(trace) = active_traces.get_mut(&span.trace_id) {
                trace.add_span(span.clone());
            }
        }
        
        // 存储活跃的Span
        {
            let mut active_spans = self.active_spans.write().await;
            active_spans.insert(span.span_id.clone(), span.clone());
        }
        
        Ok(span)
    }

    /// 从上下文开始新的Span
    pub async fn start_span_from_context(
        &self,
        context: &TracingContext,
        operation_name: &str,
        kind: SpanKind,
    ) -> ServiceResult<Span> {
        let span = if let (Some(trace_id), Some(parent_span_id)) = (&context.trace_id, &context.span_id) {
            // 创建子Span
            let mut span = Span::new_root(operation_name, &self.service_name);
            span.trace_id = trace_id.clone();
            span.parent_id = Some(parent_span_id.clone());
            span.kind = kind;
            span
        } else {
            // 创建根Span
            Span::new_root(operation_name, &self.service_name).with_kind(kind)
        };

        // 应用baggage到Span标签
        for (key, value) in &context.baggage {
            span.clone().set_tag(&format!("baggage.{}", key), value);
        }

        // 存储活跃的Span
        {
            let mut active_spans = self.active_spans.write().await;
            active_spans.insert(span.span_id.clone(), span.clone());
        }

        Ok(span)
    }

    /// 完成Span
    pub async fn finish_span(&self, span_id: &str) -> ServiceResult<()> {
        let span = {
            let mut active_spans = self.active_spans.write().await;
            active_spans.remove(span_id)
        };

        if let Some(span) = span {
            let finished_span = span.finish();
            self.processor.process_span(finished_span).await?;
        }

        Ok(())
    }

    /// 完成Trace
    pub async fn finish_trace(&self, trace_id: &str) -> ServiceResult<()> {
        let trace = {
            let mut active_traces = self.active_traces.write().await;
            active_traces.remove(trace_id)
        };

        if let Some(trace) = trace {
            let finished_trace = trace.finish();
            self.processor.process_trace(finished_trace).await?;
        }

        Ok(())
    }

    /// 获取活跃的Span
    pub async fn get_active_span(&self, span_id: &str) -> Option<Span> {
        let active_spans = self.active_spans.read().await;
        active_spans.get(span_id).cloned()
    }

    /// 更新Span
    pub async fn update_span<F>(&self, span_id: &str, updater: F) -> ServiceResult<()>
    where
        F: FnOnce(&mut Span),
    {
        let mut active_spans = self.active_spans.write().await;
        if let Some(span) = active_spans.get_mut(span_id) {
            updater(span);
        }
        Ok(())
    }

    /// 添加Span标签
    pub async fn set_span_tag(&self, span_id: &str, key: &str, value: &str) -> ServiceResult<()> {
        self.update_span(span_id, |span| {
            span.set_tag(key, value);
        }).await
    }

    /// 添加Span日志
    pub async fn log_span(&self, span_id: &str, level: &str, message: &str) -> ServiceResult<()> {
        self.update_span(span_id, |span| {
            span.log(level, message);
        }).await
    }

    /// 设置Span错误
    pub async fn set_span_error(&self, span_id: &str, error: &str) -> ServiceResult<()> {
        self.update_span(span_id, |span| {
            span.set_error(error);
        }).await
    }

    /// 设置Span状态
    pub async fn set_span_status(
        &self,
        span_id: &str,
        status: SpanStatus,
        message: Option<String>,
    ) -> ServiceResult<()> {
        self.update_span(span_id, |span| {
            span.set_status(status, message);
        }).await
    }

    /// 刷新所有待处理的追踪数据
    pub async fn flush(&self) -> ServiceResult<()> {
        self.processor.flush().await
    }

    /// 关闭追踪管理器
    pub async fn shutdown(&self) -> ServiceResult<()> {
        // 完成所有活跃的Span
        let span_ids: Vec<String> = {
            let active_spans = self.active_spans.read().await;
            active_spans.keys().cloned().collect()
        };

        for span_id in span_ids {
            self.finish_span(&span_id).await?;
        }

        // 完成所有活跃的Trace
        let trace_ids: Vec<String> = {
            let active_traces = self.active_traces.read().await;
            active_traces.keys().cloned().collect()
        };

        for trace_id in trace_ids {
            self.finish_trace(&trace_id).await?;
        }

        // 关闭处理器
        self.processor.shutdown().await
    }

    /// 获取统计信息
    pub async fn get_stats(&self) -> TracingStats {
        let active_spans_count = self.active_spans.read().await.len();
        let active_traces_count = self.active_traces.read().await.len();
        let (span_queue_size, trace_queue_size) = self.processor.get_queue_stats().await;

        TracingStats {
            active_spans_count,
            active_traces_count,
            span_queue_size,
            trace_queue_size,
        }
    }

    /// 创建手动Span（需要手动管理生命周期）
    pub async fn create_manual_span(
        &self,
        operation_name: &str,
        kind: SpanKind,
        parent_context: Option<&TracingContext>,
    ) -> ServiceResult<Span> {
        let span = if let Some(context) = parent_context {
            self.start_span_from_context(context, operation_name, kind).await?
        } else {
            self.start_root_span(operation_name).await?.with_kind(kind)
        };

        Ok(span)
    }

    /// 记录异步操作
    pub async fn trace_async<F, T>(&self, operation_name: &str, operation: F) -> ServiceResult<T>
    where
        F: std::future::Future<Output = ServiceResult<T>>,
    {
        let span = self.start_root_span(operation_name).await?;
        let span_id = span.span_id.clone();

        let start_time = SystemTime::now();
        let result = operation.await;
        let duration = start_time.elapsed().unwrap_or_default();

        // 更新Span信息
        self.update_span(&span_id, |span| {
            span.set_tag("operation.duration_ms", &duration.as_millis().to_string());
            match &result {
                Ok(_) => span.set_status(SpanStatus::Ok, None),
                Err(e) => span.set_error(&e.to_string()),
            }
        }).await?;

        self.finish_span(&span_id).await?;
        result
    }

    /// 记录同步操作
    pub async fn trace_sync<F, T>(&self, operation_name: &str, operation: F) -> ServiceResult<T>
    where
        F: FnOnce() -> ServiceResult<T>,
    {
        let span = self.start_root_span(operation_name).await?;
        let span_id = span.span_id.clone();

        let start_time = SystemTime::now();
        let result = operation();
        let duration = start_time.elapsed().unwrap_or_default();

        // 更新Span信息
        self.update_span(&span_id, |span| {
            span.set_tag("operation.duration_ms", &duration.as_millis().to_string());
            match &result {
                Ok(_) => span.set_status(SpanStatus::Ok, None),
                Err(e) => span.set_error(&e.to_string()),
            }
        }).await?;

        self.finish_span(&span_id).await?;
        result
    }
}

/// 追踪统计信息
#[derive(Debug, Clone)]
pub struct TracingStats {
    /// 活跃Span数量
    pub active_spans_count: usize,
    /// 活跃Trace数量
    pub active_traces_count: usize,
    /// Span队列大小
    pub span_queue_size: usize,
    /// Trace队列大小
    pub trace_queue_size: usize,
}

/// 追踪装饰器宏辅助结构
pub struct TracedOperation<'a> {
    manager: &'a TracingManager,
    span_id: String,
}

impl<'a> TracedOperation<'a> {
    pub fn new(manager: &'a TracingManager, span_id: String) -> Self {
        Self { manager, span_id }
    }

    /// 添加标签
    pub async fn tag(&self, key: &str, value: &str) -> ServiceResult<()> {
        self.manager.set_span_tag(&self.span_id, key, value).await
    }

    /// 添加日志
    pub async fn log(&self, level: &str, message: &str) -> ServiceResult<()> {
        self.manager.log_span(&self.span_id, level, message).await
    }

    /// 设置错误
    pub async fn error(&self, error: &str) -> ServiceResult<()> {
        self.manager.set_span_error(&self.span_id, error).await
    }
}

impl<'a> Drop for TracedOperation<'a> {
    fn drop(&mut self) {
        // 注意：这里不能使用async，所以需要在外部确保Span被正确完成
        // 实际使用时应该通过try块或其他方式确保完成
    }
}

/// 简化的追踪API
pub struct SimpleTracer {
    manager: Arc<TracingManager>,
}

impl SimpleTracer {
    pub fn new(manager: Arc<TracingManager>) -> Self {
        Self { manager }
    }

    /// 开始新的操作追踪
    pub async fn start(&self, operation_name: &str) -> ServiceResult<String> {
        let span = self.manager.start_root_span(operation_name).await?;
        Ok(span.span_id)
    }

    /// 完成操作追踪
    pub async fn finish(&self, span_id: &str) -> ServiceResult<()> {
        self.manager.finish_span(span_id).await
    }

    /// 记录事件
    pub async fn event(&self, span_id: &str, message: &str) -> ServiceResult<()> {
        self.manager.log_span(span_id, "info", message).await
    }

    /// 记录错误
    pub async fn error(&self, span_id: &str, error: &str) -> ServiceResult<()> {
        self.manager.set_span_error(span_id, error).await
    }

    /// 添加属性
    pub async fn attribute(&self, span_id: &str, key: &str, value: &str) -> ServiceResult<()> {
        self.manager.set_span_tag(span_id, key, value).await
    }
}